#!/usr/bin/env python
# -*- coding:utf-8 -*-
# Copyright 2022-2022 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "license");
# you may not use this file except in compliance with the License.
# you may obtain a copy of the License at
#
# http://www.apache.org/license/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==========================================================================
"""
compile operator
"""

import os
import stat
import json
import platform
import subprocess
from tbe.tvm.contrib.ccec import _build_aicore_compile_cmd, run_cmd
from tbe.tvm.runtime.cce_runtime import tvm_callback_cce_postproc
from tbe.common.buildcfg import get_current_build_config
from tbe.tvm.error_mgr import raise_tbe_python_err, TBE_DEFAULT_PYTHON_ERROR_CODE
from tbe.common.utils.op_tiling import do_op_tiling
from tbe.tikcpp.get_op_tiling import load_lib, get_tiling_def, get_tiling_data, gen_static_shape, gen_dynamic_shape
from te.tik.common.tik_get_soc_name import get_soc_name

BISHENG_CPP_HOME = os.environ.get('BISHENG_CPP_HOME')
if BISHENG_CPP_HOME  is None:
    raise RuntimeError(f"env BISHENG_CPP_HOME not exists, please export BISHENG_CPP_HOME.")
if not os.path.exists(BISHENG_CPP_HOME):
    raise RuntimeError(f"dir {BISHENG_CPP_HOME} not exists, please export BISHENG_CPP_HOME.")

SOC_NAME = get_soc_name()
if SOC_NAME.find('310') != -1:
    BSCPP_SYCL_TARGET = "ascend_310-cce"
elif SOC_NAME.find('910') != -1:
    BSCPP_SYCL_TARGET = "ascend_910-cce"
else:
    BSCPP_SYCL_TARGET = "ascend_910-cce"

CANN_ARCH = platform.machine() + "-linux"


def compile_op(bscpp_src_file: str, kernel_name: str, optype: str, inputs: list, outputs: list):
    """get tiling_data/ generate tiling_data file/ compile cce to .o / generate .json file

    Args:
        bscpp_src_file (str):BiShengCPP file to be compiled
        kernel_name (str): kernel name, to generate .o and .json file
        optype (str): operator type
        inputs (list): shape/ori_shape/format/ori_format/dtype/name, type including required/optional/dynamic,
                        if shape of inputs is const, there will be an extra field "const value"
        outputs (list): shape/ori_shape/format/ori_format/dtype/name, type including required/optional/dynamic
    """
    # The registration of the tiling is done in do_op_tiling
    # check op_tiling.py to see if we can call or refuse the api
    # if static shape, call TbeOpTilingPyInterface, return tilingdata
    # elif dynamic shape, do nothing
    print(f"[INFO] bscpp.compile_op bscpp_src_file:{bscpp_src_file}, kernel_name:{kernel_name}, optype:{optype}")
    print(get_soc_name())
    run_info = {}
    load_lib()

    tiling_def = get_tiling_def(optype)
    if tiling_def is None:
        run_info["block_dim"] = -1
    else:
        static_shape = _is_static_shape(inputs)
        if static_shape:
            compile_info = ""
            run_info = do_op_tiling(optype, compile_info, inputs, outputs)
            run_info = get_tiling_data(tiling_def, run_info)
        else:
            run_info["block_dim"] = -1
    tiling_data_num = 16
    kernel_meta_dir = _get_kernel_meta_dir()
    dest_kernel_file = os.path.join(kernel_meta_dir, kernel_name + ".o")
    tiling_data_size = 4 * tiling_data_num
    try:
        block_dim = run_info["block_dim"]
        kernel_attrs = {"kernel_name": kernel_name, "block_dim": block_dim, "arg_size": tiling_data_size}
        compile_bscpp_cce(bscpp_src_file, dest_kernel_file, kernel_attrs)
    except KeyError:
        print("[ERROR] block_dim is not in run_info， run compile_bscpp_cce failed")
        sys.exit(1)


def compile_bscpp_cce(bscpp_src_file: str, dest_kernel_file: str, kernel_attrs: dict):
    """call BiShengCPP to compile a BiShengCPP source file, generate a kernel binary file and a json file

    Args:
        bscpp_src_file (str):bisheng file
        dest_kernel_file (str): bisheng kernel binary file
        kernel_attrs (dict): block_dim/arg_size/kernel_name
    """
    out_file = dest_kernel_file.replace(".o", ".out")
    cmd = f"{BISHENG_CPP_HOME}/bin/clang++ -fsycl -fsycl-targets={BSCPP_SYCL_TARGET} {bscpp_src_file} \
        --cce-o={dest_kernel_file} -o {out_file}"
    command = cmd.strip().split(' ')
    print(f"[INFO] bscpp build source start, cmd:{cmd}")
    ret = subprocess.run(command, stderr=subprocess.PIPE)
    if ret.returncode == 0:
        print("[INFO] bscpp build source success: ", ret)
    else:
        print("[ERROR] bscpp build source filed: ", ret)

    target = "cce_core"
    kernel_name = kernel_attrs["kernel_name"]
    block_dim = kernel_attrs["block_dim"]
    arg_size = kernel_attrs["arg_size"]
    tvm_callback_cce_postproc(target, kernel_name, block_dim, arg_size)
    kernel_func_name = _get_kernel_binary_kernel_func_name(dest_kernel_file)
    if kernel_func_name != "":
        _modify_kernel_json_kernel_func_name(dest_kernel_file, kernel_func_name)


def _modify_kernel_json_kernel_func_name(dest_kernel_file, kernel_func_name):
    print("[INFO] kernelName: ", kernel_func_name)
    kernel_json = dest_kernel_file.replace(".o", ".json")
    with open(kernel_json, 'r') as f:
        json_obj = json.load(f)
    json_obj["kernelName"] = kernel_func_name

    fd = os.open(kernel_json, os.O_WRONLY | os.O_TRUNC, stat.S_IWOTH | stat.S_IWGRP | stat.S_IWUSR)
    with os.fdopen(fd, 'w') as f:
        json.dump(json_obj, f, indent=4, ensure_ascii=False)
    print("[INFO] replace kernel name in json success")


def _get_kernel_binary_kernel_func_name(dest_kernel_file):
    cmd = f"objdump -t {dest_kernel_file} |grep KernelName"
    out = subprocess.getoutput(cmd)
    lines = out.split("\n")
    for line in lines:
        fields = line.split(" ")
        if len(fields) < 3:
            continue
        kernel_func_name = fields[len(fields) - 1]
        if kernel_func_name.endswith("$local"):
            continue
        return kernel_func_name
    return ""


def _is_static_shape(inputs: list):
    """check if static shape, find dynamic shape if shape<0 in inputs

    Args:
    inputs (list): shape/ori_shape/format/ori_format/dtype/name, type including required/optional/dynamic,
                    if shape of inputs is const, there will be an extra field "const value".
    Returns:
        res (Boolean): True means static_shape, False means dynamic shape
    """
    for input_ele in inputs:
        for shape in input_ele["shape"]:
            if shape < 0:
                return False
    return True


def _get_kernel_meta_dir():
    kernel_meta_dir = os.path.join(get_current_build_config("kernel_meta_parent_dir"), "kernel_meta")
    if not os.path.exists(kernel_meta_dir):
        try:
            os.makedirs(kernel_meta_dir, stat.S_IRWXU + stat.S_IRGRP + stat.S_IXGRP, exist_ok=True)
        except OSError as err:
            print(f"create {kernel_meta_dir} fail")
            raise_tbe_python_err(TBE_DEFAULT_PYTHON_ERROR_CODE, err)
    else:
        print(f"{kernel_meta_dir} already exists")
        ret = subprocess.run(["ls", "{kernel_meta_dir}"], stderr=subprocess.PIPE)
        if ret.returncode != 0:
            print("ERROR: ", ret)
    return kernel_meta_dir

